//
// Created by LIN on 25-8-20.
//

#ifndef RBTREE_H
#define RBTREE_H
#include <print>

using namespace std;

// 枚举值表示颜色
enum Colour
{
    red,
    black
};

// 这里我们默认按key/value结构实现
template<class K, class V>
struct RBTreeNode
{
    // 这里更新控制平衡也要加入parent指针
    pair<K, V> kv_;
    RBTreeNode* left_;
    RBTreeNode* right_;
    RBTreeNode* parent_;
    Colour col_;
    explicit RBTreeNode(const pair<K, V>& kv)
        : kv_(kv)
        , left_(nullptr)
        , right_(nullptr)
        , parent_(nullptr)
        , col_(black)
    {}
};

template<class K, class V>
class RBTree
{
    using Node = RBTreeNode<K, V>;

public:
    bool insert(const pair<K, V>& kv)
    {
        if (root_ == nullptr)
        {
            root_ = new Node(kv);
            root_->col_ = black;
        }

        Node* parent = nullptr;
        Node* current = root_;

        while (current)
        {
            if (current->kv_.first < kv.first)
            {
                parent = current;
                current = current->right_;
            }
            else if (current->kv_.first > kv.first)
            {
                parent = current;
                current = current->left_;
            }
            else
            {
                return false;
            }
        }

        current = new Node(kv);
        current->col_ = red;

        if (parent->kv_.first < parent->kv_.first)
        {
            parent->right_ = current;
        }
        else
        {
            parent->left_ = current;
        }

        current->parent_ = parent;

        while (parent && parent->col_ == red)
        {
            Node* grandfather = parent->parent_;
            if (grandfather->left_ == parent)
            {
                Node* uncle = grandfather->right_;
                if (uncle && uncle->col_ == red)
                {
                    parent->col_ = black;
                    uncle->col_ = black;
                    grandfather->col_ = red;
                }
            }
        }
    }

private:
    Node* root_ = nullptr;
};

#endif /* RBTREE_H */
